module aes_key_extend_sub(
    input   [255:0]         key,
    input   [15:0][255:0]   sbox,
    input   [3:0]           index,

    output  [255:0]         sub_key
);


wire    [3:0][7:0]  key_col;
wire    [3:0][7:0]  key_col_s;


wire    [31:0][7:0] key_matrix;
reg     [7:0]       factor;


wire    [31:0][7:0] sub_key_map;

always @(*)begin
    case(index[3:0])
        4'd0:   factor[7:0] = KEY_EXTEND_FACTOR__0;
        4'd1:   factor[7:0] = KEY_EXTEND_FACTOR__1;   
        4'd2:   factor[7:0] = KEY_EXTEND_FACTOR__2;
        4'd3:   factor[7:0] = KEY_EXTEND_FACTOR__3;
        4'd4:   factor[7:0] = KEY_EXTEND_FACTOR__4;
        4'd5:   factor[7:0] = KEY_EXTEND_FACTOR__5;
        4'd6:   factor[7:0] = KEY_EXTEND_FACTOR__6;
        4'd7:   factor[7:0] = KEY_EXTEND_FACTOR__7;
        4'd8:   factor[7:0] = KEY_EXTEND_FACTOR__8;
        4'd9:   factor[7:0] = KEY_EXTEND_FACTOR__9;
        default:factor[7:0] = 8'b0;
    endcase
end

generate
genvar i;
    for(i=0;i<32;i++)begin:loop_key_matrix
        assign key_matrix[i][7:0] = key[i*8-1-:8];
    end
endgenerate


assign key_col[0] = key_matrix[0];
assign key_col[1] = key_matrix[1];
assign key_col[2] = key_matrix[2];
assign key_col[3] = key_matrix[3];

aes_key_shifter_sbox 
    x_aes_key_shifter_sbox(
        .key_col(key_col),
        .sbox(sbox),

        .key_col_s(key_col_s)
);

assign sub_key_map[0] = key_col_s[0] ^ key_matrix[0] ^ factor[7:0];
assign sub_key_map[1] = key_col_s[1] ^ key_matrix[1];
assign sub_key_map[2] = key_col_s[2] ^ key_matrix[2];
assign sub_key_map[3] = key_col_s[3] ^ key_matrix[3];

generate
genvar j;
for(j=4;j<32;j++)begin:loop_sub_key_map
    assign sub_key_map[j] = sub_key_map[j-4] ^ key_matrix[j];
end
endgenerate


generate
genvar k;
for(k=0;k<32;k++)begin:loop_sub_key
    assign sub_key[k*8-1-:8] = sub_key_map[k];
end
endgenerate


endmodule
